from docx import Document
import os
import json


class TableStructureValidator:
    def __init__(self, doc_zh_path, doc_en_path, checkpoint_file='checkpoint.json'):
        self.doc_zh_path = doc_zh_path
        self.doc_en_path = doc_en_path
        self.checkpoint_file = checkpoint_file
        self.current_index = 0
        self.total_tables = 0
        self.mismatch_log = []

        # 加载上次的校验进度（如果有）
        self.load_checkpoint()

    def load_checkpoint(self):
        """加载上次的校验进度"""
        if os.path.exists(self.checkpoint_file):
            try:
                with open(self.checkpoint_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                # self.current_index = data.get('current_index', 0)
                self.current_index =  0
                self.total_tables = data.get('total_tables', 0)
                self.mismatch_log = data.get('mismatch_log', [])
                print(f"✓ 已加载校验进度: 从表格 {self.current_index + 1} 开始")
            except Exception as e:
                print(f"⚠ 检查点文件损坏，将从头开始校验: {str(e)}")

    def save_checkpoint(self):
        """保存当前校验进度"""
        data = {
            'current_index': self.current_index,
            'total_tables': self.total_tables,
            'mismatch_log': self.mismatch_log,
            'doc_zh_path': self.doc_zh_path,
            'doc_en_path': self.doc_en_path
        }
        with open(self.checkpoint_file, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    def extract_table_structure(self, table):
        """
        提取表格结构特征
        :param table: docx表格对象
        :return: 结构特征字典
        """
        # 基本维度
        rows = len(table.rows)
        cols = len(table.columns) if table.rows and hasattr(table, 'columns') else 0

        # 检测合并单元格模式
        merge_pattern = []
        for row_idx, row in enumerate(table.rows):
            # 获取行中的单元格列表
            cells = row.cells
            for cell_idx, cell in enumerate(cells):
                # 检测横向合并
                col_span = 1
                if hasattr(cell, '_tc') and hasattr(cell._tc, 'tcPr') and cell._tc.tcPr is not None and hasattr(
                        cell._tc.tcPr, 'gridSpan') and cell._tc.tcPr.gridSpan is not None:
                    col_span = cell._tc.tcPr.gridSpan.val

                # 检测纵向合并
                row_span = 1
                if hasattr(cell, '_tc') and hasattr(cell._tc, 'tcPr') and cell._tc.tcPr is not None and hasattr(
                        cell._tc.tcPr,
                        'vMerge') and cell._tc.tcPr.vMerge is not None and cell._tc.tcPr.vMerge.val == 'restart':
                    row_span = 2  # 简化处理：标记有纵向合并

                if col_span > 1 or row_span > 1:
                    merge_pattern.append({
                        'position': (row_idx, cell_idx),
                        'col_span': col_span,
                        'row_span': row_span
                    })

        return {
            'dimensions': (rows, cols),
            'merge_pattern': merge_pattern
        }

    def extract_table_content(self, table, max_rows=10, max_cols=10):
        """
        提取表格内容（前max_rows行和前max_cols列）
        :param table: docx表格对象
        :param max_rows: 最大提取行数
        :param max_cols: 最大提取列数
        :return: 表格内容列表
        """
        content = []
        for row_idx, row in enumerate(table.rows):
            if row_idx >= max_rows:
                break
            row_content = []
            # 获取行中的单元格列表
            cells = row.cells
            for cell_idx, cell in enumerate(cells):
                if cell_idx >= max_cols:
                    break
                # 提取单元格文本
                cell_text = cell.text.strip()
                row_content.append(cell_text)
            content.append(row_content)
        return content

    def compare_table_structures(self, struct_zh, struct_en):
        """
        比较两个表格结构
        :return: (是否匹配, 差异描述)
        """
        # 1. 检查维度
        rows_zh, cols_zh = struct_zh['dimensions']
        rows_en, cols_en = struct_en['dimensions']

        if rows_zh != rows_en:
            return False, f"行数不同: {rows_zh} (中) vs {rows_en} (英)"

        if cols_zh != cols_en:
            return False, f"列数不同: {cols_zh} (中) vs {cols_en} (英)"

        # 2. 检查合并单元格模式
        zh_merge_set = set((m['position'], m['col_span'], m['row_span']) for m in struct_zh['merge_pattern'])
        en_merge_set = set((m['position'], m['col_span'], m['row_span']) for m in struct_en['merge_pattern'])

        if zh_merge_set != en_merge_set:
            # 找出具体差异
            missing_in_en = zh_merge_set - en_merge_set
            extra_in_en = en_merge_set - zh_merge_set

            diff_msg = "合并单元格模式不同:"
            if missing_in_en:
                diff_msg += f"\n  - 中文有但英文无: {missing_in_en}"
            if extra_in_en:
                diff_msg += f"\n  - 英文有但中文无: {extra_in_en}"
            return False, diff_msg

        return True, "结构匹配"

    def validate_tables(self):
        """执行表格结构校验"""
        print(f"\n{'=' * 50}")
        print(f"开始校验文档表格结构")
        print(f"中文文档: {self.doc_zh_path}")
        print(f"英文文档: {self.doc_en_path}")
        print(f"从表格 {self.current_index + 1} 开始校验")
        print(f"{'=' * 50}\n")

        # 加载文档
        try:
            doc_zh = Document(self.doc_zh_path)
            doc_en = Document(self.doc_en_path)
        except Exception as e:
            print(f"⚠ 加载文档失败: {str(e)}")
            return False

        tables_zh = doc_zh.tables
        tables_en = doc_en.tables

        # 检查表格总数
        if self.total_tables == 0:
            self.total_tables = min(len(tables_zh), len(tables_en))

        if len(tables_zh) != len(tables_en):
            print(f"⚠ 警告: 文档表格数量不同 - 中文: {len(tables_zh)} 英文: {len(tables_en)}")
            print("将校验两个文档共有的表格数量")

        # 顺序校验表格
        for idx in range(self.current_index, self.total_tables):
            print(f"\n🔍 校验表格 {idx + 1}/{self.total_tables}...")

            try:
                table_zh = tables_zh[idx]
                table_en = tables_en[idx]
            except IndexError:
                print(f"⚠ 表格 {idx + 1} 不存在于其中一个文档中，跳过")
                self.current_index = idx + 1
                self.save_checkpoint()
                continue

            # 提取结构
            try:
                struct_zh = self.extract_table_structure(table_zh)
                struct_en = self.extract_table_structure(table_en)
            except Exception as e:
                print(f"⚠ 提取表格结构失败: {str(e)}")
                self.current_index = idx
                self.save_checkpoint()
                return False

            # 比较结构
            is_match, diff_msg = self.compare_table_structures(struct_zh, struct_en)

            if is_match:
                print(f"  ✓ 表格 {idx + 1} 结构匹配")
                self.current_index = idx + 1
                self.save_checkpoint()
            else:
                print(f"  ✗ 表格 {idx + 1} 结构不匹配")
                print(f"    差异: {diff_msg}")

                # 提取表格内容用于显示
                try:
                    content_zh = self.extract_table_content(table_zh)
                    content_en = self.extract_table_content(table_en)
                except Exception as e:
                    print(f"⚠ 提取表格内容失败: {str(e)}")
                    content_zh = []
                    content_en = []

                # 记录不匹配
                self.mismatch_log.append({
                    'table_index': idx,
                    'diff_message': diff_msg,
                    'zh_structure': struct_zh,
                    'en_structure': struct_en,
                    'zh_content': content_zh,
                    'en_content': content_en
                })

                # 保存检查点
                self.current_index = idx
                self.save_checkpoint()

                # 输出摘要
                self.print_summary()
                return False

        # 所有表格校验完成
        self.current_index = self.total_tables
        self.save_checkpoint()

        if self.mismatch_log:
            print("\n⚠ 警告: 所有表格已完成校验，但存在不匹配的记录")
            self.print_summary()
            return False
        else:
            print("\n🎉 所有表格结构匹配!")
            # 清除检查点文件
            if os.path.exists(self.checkpoint_file):
                os.remove(self.checkpoint_file)
            return True

    def format_table_content(self, content, title, max_rows=10, max_cols=10):
        """
        格式化表格内容用于显示
        :param content: 表格内容列表
        :param title: 表格标题
        :param max_rows: 最大显示行数
        :param max_cols: 最大显示列数
        :return: 格式化后的字符串
        """
        if not content:
            return f"\n{title}:\n(无内容或提取失败)\n"

        formatted = f"\n{title}:\n"
        formatted += "-" * 60 + "\n"

        # 添加列号标题
        # 找出所有行中的最大列数
        max_cols_in_content = max(len(row) for row in content) if content else 0
        cols_to_show = min(max_cols, max_cols_in_content)

        col_headers = [f"列{i + 1}" for i in range(cols_to_show)]
        formatted += "行号 | " + " | ".join(col_headers) + "\n"
        formatted += "-" * 60 + "\n"

        # 添加行内容
        for row_idx, row in enumerate(content[:max_rows]):
            # 确保该行有足够的列
            row_display = []
            for col_idx in range(cols_to_show):
                if col_idx < len(row):
                    cell_text = row[col_idx]
                    # 截断长文本
                    if len(cell_text) > 50:
                        cell_text = cell_text[:47] + "..."
                    row_display.append(cell_text)
                else:
                    row_display.append("")

            formatted += f"{row_idx + 1:3} | " + " | ".join(f"{cell:<20}" for cell in row_display) + "\n"

        # 添加截断提示
        if len(content) > max_rows:
            formatted += f"... 还有 {len(content) - max_rows} 行未显示\n"
        if max_cols_in_content > max_cols:
            formatted += f"... 还有 {max_cols_in_content - max_cols} 列未显示\n"

        formatted += "-" * 60
        return formatted

    def print_summary(self):
        """输出校验摘要，包含表格内容对比"""
        print(f"\n{'=' * 50}")
        print("表格结构校验摘要")
        print(f"{'=' * 50}")
        print(f"已校验表格: {self.current_index}/{self.total_tables}")
        print(f"不匹配表格数: {len(self.mismatch_log)}")

        if not self.mismatch_log:
            return

        last_mismatch = self.mismatch_log[-1]
        table_idx = last_mismatch['table_index'] + 1
        print(f"\n最后不匹配的表格: {table_idx}")
        print(f"差异原因: {last_mismatch['diff_message']}")

        zh_rows, zh_cols = last_mismatch['zh_structure']['dimensions']
        en_rows, en_cols = last_mismatch['en_structure']['dimensions']
        print(f"\n中文表格维度: {zh_rows}行 × {zh_cols}列")
        print(f"英文表格维度: {en_rows}行 × {en_cols}列")

        # 显示合并单元格差异
        zh_merge = last_mismatch['zh_structure']['merge_pattern']
        en_merge = last_mismatch['en_structure']['merge_pattern']
        print(f"\n中文合并单元格数: {len(zh_merge)}")
        print(f"英文合并单元格数: {len(en_merge)}")

        if zh_merge or en_merge:
            print("\n合并单元格详情:")
            print("位置格式: (行, 列) [跨列数×跨行数]")

            if zh_merge:
                print(" 中文:")
                for m in zh_merge:
                    pos = m['position']
                    print(f"  - {pos} [{m['col_span']}×{m['row_span']}]")

            if en_merge:
                print(" 英文:")
                for m in en_merge:
                    pos = m['position']
                    print(f"  - {pos} [{m['col_span']}×{m['row_span']}]")

        # 输出表格内容对比
        print(f"\n{'=' * 50}")
        print(f"表格 {table_idx} 内容对比 (前10行前10列)")
        print(f"{'=' * 50}")

        # 格式化表格内容
        zh_content = self.format_table_content(
            last_mismatch['zh_content'],
            f"中文表格 {table_idx} 内容"
        )
        en_content = self.format_table_content(
            last_mismatch['en_content'],
            f"英文表格 {table_idx} 内容"
        )

        print(zh_content)
        print(en_content)

        # 高亮显示内容差异
        print("\n内容差异高亮 (前10行前10列):")
        found_difference = False

        # 确定要对比的行数和列数
        rows_to_compare = min(10, len(last_mismatch['zh_content']), len(last_mismatch['en_content']))
        cols_to_compare = min(10, max(
            max(len(row) for row in last_mismatch['zh_content'][:rows_to_compare]) if last_mismatch[
                'zh_content'] else 0,
            max(len(row) for row in last_mismatch['en_content'][:rows_to_compare]) if last_mismatch['en_content'] else 0
        ))

        for row_idx in range(rows_to_compare):
            # 确保行索引有效
            zh_row = last_mismatch['zh_content'][row_idx] if row_idx < len(last_mismatch['zh_content']) else []
            en_row = last_mismatch['en_content'][row_idx] if row_idx < len(last_mismatch['en_content']) else []

            for col_idx in range(cols_to_compare):
                # 确保列索引有效
                zh_cell = zh_row[col_idx] if col_idx < len(zh_row) else ""
                en_cell = en_row[col_idx] if col_idx < len(en_row) else ""

                # 比较单元格内容
                if zh_cell != en_cell:
                    found_difference = True
                    print(f"行 {row_idx + 1}, 列 {col_idx + 1}:")
                    if zh_cell:
                        print(f"  中文: {zh_cell[:100]}{'...' if len(zh_cell) > 100 else ''}")
                    else:
                        print(f"  中文: (空)")
                    if en_cell:
                        print(f"  英文: {en_cell[:100]}{'...' if len(en_cell) > 100 else ''}")
                    else:
                        print(f"  英文: (空)")
                    print("")

        if not found_difference:
            print("  前10行前10列内容完全相同（但表格结构不匹配，请检查合并单元格或行列数）")

        print(f"\n操作说明:")
        print(f"1. 请修改文档中的表格 {self.current_index + 1}")
        print(f"2. 保存文档后重新运行此程序继续校验")
        print(f"3. 校验进度已保存到: {self.checkpoint_file}")
        print(f"{'=' * 50}")


def main():
    # 配置文档路径
    doc_zh = "./test/历史、重组及公司架构/5 中文招股书——历史、重组及公司架构.docx"  # 中文文档路径
    doc_en = "./test/历史、重组及公司架构/5 英文招股书——历史、重组及公司架构.docx"  # 英文文档路径
    checkpoint = "table_validation_checkpoint.json"

    # 检查文件是否存在
    if not os.path.exists(doc_zh):
        print(f"错误：中文文档不存在 - {doc_zh}")
        return
    if not os.path.exists(doc_en):
        print(f"错误：英文文档不存在 - {doc_en}")
        return

    # 创建校验器
    validator = TableStructureValidator(doc_zh, doc_en, checkpoint)

    # 执行校验
    is_complete = validator.validate_tables()

    # 输出最终结果
    if is_complete:
        print("\n✅ 所有表格校验完成且完全匹配!")
    else:
        print("\n❌ 校验暂停，请修改文档后重新运行")


if __name__ == "__main__":
    main()